"""
顺序工作流代理示例 - 智能文档处理系统
演示如何使用SequentialAgent按顺序执行多个代理来处理文档
"""

from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.tools.function_tool import FunctionTool
from google.adk.tools.tool_context import ToolContext
import os
from google.adk.models.lite_llm import LiteLlm

# API密钥配置
DEFAULT_DASHSCOPE_API_KEY = "sk-f227634bb56*************************232"  # 请替换为您的实际密钥
DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY", DEFAULT_DASHSCOPE_API_KEY)

# 常量定义
APP_NAME = "sequential_document_processor"
USER_ID = "dev_user_01"
GEMINI_MODEL = LiteLlm(
    model="openai/qwen-turbo",  # 使用通义千问Turbo模型
    api_key=DASHSCOPE_API_KEY,
    api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
)

# 状态键定义
STATE_RAW_DOCUMENT = "raw_document"
STATE_EXTRACTED_CONTENT = "extracted_content"
STATE_ANALYSIS_RESULT = "analysis_result"
STATE_SUMMARY = "summary"
STATE_FINAL_OUTPUT = "final_output"

# 示例文档内容
SAMPLE_DOCUMENT = """
人工智能技术发展报告

随着深度学习和大语言模型的快速发展，人工智能技术正在各个领域产生深远影响。
本报告分析了当前AI技术的主要趋势和应用场景。

主要发现：
1. 生成式AI在内容创作领域表现突出
2. 多模态AI模型能够处理文本、图像、音频等多种数据类型
3. AI代理系统在自动化任务执行方面展现巨大潜力

技术挑战：
- 模型的可解释性和透明度
- 数据隐私和安全问题
- 计算资源的高消耗

未来展望：
AI技术将继续向更加智能化、个性化的方向发展，预计在教育、医疗、金融等领域
将有更多突破性应用。同时，需要建立完善的AI治理框架来确保技术的负责任发展。

结论：
人工智能技术正处于快速发展期，机遇与挑战并存。企业和研究机构需要在技术创新
和风险管控之间找到平衡点，推动AI技术的健康发展。
"""

# 工具函数定义
def extract_key_information(tool_context: ToolContext):
    """从原始文档中提取关键信息的工具"""
    print(f"[Tool] extract_key_information called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 从用户输入或会话状态中获取文档内容
    raw_doc = session_state.get(STATE_RAW_DOCUMENT, "")
    
    # 如果状态中没有文档，尝试从用户输入中获取
    if not raw_doc:
        # 尝试从用户内容中获取文本
        user_content = tool_context._invocation_context.user_content
        if user_content and hasattr(user_content, 'parts'):
            # 从Content对象的parts中提取文本
            for part in user_content.parts:
                if hasattr(part, 'text') and part.text:
                    raw_doc = part.text
                    break
        elif user_content:
            raw_doc = str(user_content)
    
    # 如果还是没有，使用示例文档
    if not raw_doc:
        raw_doc = SAMPLE_DOCUMENT
    
    # 保存到状态中
    session_state[STATE_RAW_DOCUMENT] = raw_doc
    
    # 这里可以集成真实的文档解析库，如PyPDF2, python-docx等
    extracted_info = {
        "title": "从文档中提取的标题",
        "key_points": ["要点1", "要点2", "要点3"],
        "entities": ["实体1", "实体2"],
        "word_count": len(raw_doc.split()) if raw_doc else 0,
        "document_preview": raw_doc[:200] + "..." if len(raw_doc) > 200 else raw_doc
    }
    
    session_state[STATE_EXTRACTED_CONTENT] = extracted_info
    return extracted_info

def format_final_document(tool_context: ToolContext):
    """格式化最终文档的工具"""
    print(f"[Tool] format_final_document called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    summary = session_state.get(STATE_SUMMARY, "")
    analysis = session_state.get(STATE_ANALYSIS_RESULT, "")
    
    formatted_output = f"""
# 文档处理报告

## 摘要
{summary}

## 详细分析
{analysis}

## 处理时间
{session_state.get('processing_timestamp', 'N/A')}

---
*本报告由ADK顺序工作流代理自动生成*
"""
    
    session_state[STATE_FINAL_OUTPUT] = formatted_output
    
    # 设置结束调用标志
    tool_context._invocation_context.end_invocation = True
    
    # 返回明确的完成信号
    return {
        "status": "SUCCESS",
        "message": "文档格式化已完成，任务结束",
        "formatted_document": formatted_output,
        "task_completed": True
    }

# 代理定义
# 转换为ADK工具格式
extract_key_information_tool = FunctionTool(extract_key_information)
# 步骤1: 内容提取代理
content_extractor_agent = LlmAgent(
    name="ContentExtractorAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个专业的文档内容提取专家。
    
    你的任务是分析用户提供的文档内容，并使用extract_key_information工具提取关键信息。
    
    用户输入的内容将作为文档进行处理。请调用extract_key_information工具来提取文档的关键信息，包括：
    - 标题和主题
    - 关键要点
    - 重要实体
    - 文档统计信息
    
    提取完成后，简要说明提取的内容类型和数量。
    """,
    description="负责从原始文档中提取关键信息和结构化数据",
    tools=[extract_key_information_tool],
    include_contents='none'
)

# 步骤2: 内容分析代理
content_analyzer_agent = LlmAgent(
    name="ContentAnalyzerAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个专业的内容分析师。
    
    基于之前提取的内容信息，请进行深入分析。
    
    请分析以下方面：
    1. 文档的主要主题和核心观点
    2. 内容的逻辑结构和组织方式
    3. 关键信息的重要性排序
    4. 文档的目标受众和用途
    5. 内容的完整性和质量评估
    
    请提供结构化的分析结果，包括具体的观察和建议。
    """,
    description="分析提取的内容，识别主题、结构和关键点",
    output_key=STATE_ANALYSIS_RESULT,
    include_contents='none'
)

# 步骤3: 摘要生成代理
summary_generator_agent = LlmAgent(
    name="SummaryGeneratorAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个专业的文档摘要生成专家。
    
    基于之前的内容分析结果，生成一个简洁而全面的文档摘要。
    
    请生成一个包含以下要素的摘要：
    1. 文档核心主题（1-2句话）
    2. 主要观点和结论（3-5个要点）
    3. 关键数据或事实
    4. 实用价值和应用建议
    
    摘要应该：
    - 长度控制在200-300字
    - 语言简洁明了
    - 突出重点信息
    - 便于快速理解
    """,
    description="基于分析结果生成结构化的文档摘要",
    output_key=STATE_SUMMARY,
    include_contents='none'
)

# 转换为ADK工具格式
format_final_document_tool = FunctionTool(format_final_document)
# 步骤4: 格式化代理
formatter_agent = LlmAgent(
    name="FormatterAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个专业的文档格式化专家。

    你的任务是调用format_final_document工具一次来格式化最终输出文档，然后立即结束。

    工作流程：
    1. 调用format_final_document工具
    2. 确认工具返回成功状态
    3. 立即结束，不要再次调用任何工具

    重要：只调用工具一次，收到工具响应后立即结束任务。
    """,
    description="将所有处理结果格式化为最终的输出文档",
    tools=[format_final_document_tool],
    include_contents='none'
)

# 创建顺序工作流代理
root_agent = SequentialAgent(
    name=APP_NAME,
    description="智能文档处理系统 - 按顺序执行内容提取、分析、摘要生成和格式化",
    sub_agents=[
        content_extractor_agent,
        content_analyzer_agent,
        summary_generator_agent,
        formatter_agent
    ]
)

# 用于测试的辅助函数
def create_test_session_with_document(session_service, document_content: str):
    """创建包含测试文档的会话"""
    session = session_service.create_session(APP_NAME, USER_ID)
    session.state[STATE_RAW_DOCUMENT] = document_content
    session.state['processing_timestamp'] = "2024-01-15 10:30:00"
    return session